/*
    Copyright (C) 2012-2013 de4dot@gmail.com

    Permission is hereby granted, free of charge, to any person obtaining
    a copy of this software and associated documentation files (the
    "Software"), to deal in the Software without restriction, including
    without limitation the rights to use, copy, modify, merge, publish,
    distribute, sublicense, and/or sell copies of the Software, and to
    permit persons to whom the Software is furnished to do so, subject to
    the following conditions:

    The above copyright notice and this permission notice shall be
    included in all copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
    IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
    CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
    TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
    SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

﻿using System;
using System.IO;
using System.Security.Cryptography;

namespace dnlib.DotNet {
	/// <summary>
	/// Thrown if the strong name key or public key is invalid
	/// </summary>
	[Serializable]
	public class InvalidKeyException : Exception {
		/// <summary>
		/// Default constructor
		/// </summary>
		public InvalidKeyException() {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="message">Error message</param>
		public InvalidKeyException(string message)
			: base(message) {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="message">Error message</param>
		/// <param name="innerException">Other exception</param>
		public InvalidKeyException(string message, Exception innerException)
			: base(message, innerException) {
		}
	}

	/// <summary>
	/// Type of signature algorithm. See WinCrypt.h in the Windows SDK
	/// </summary>
	public enum SignatureAlgorithm : uint {
		/// <summary>
		/// RSA signature algorithm
		/// </summary>
		CALG_RSA_SIGN = 0x00002400,
	}

	static class StrongNameUtils {
		public static byte[] ReadBytesReverse(this BinaryReader reader, int len) {
			var data = reader.ReadBytes(len);
			if (data.Length != len)
				throw new InvalidKeyException("Can't read more bytes");
			Array.Reverse(data);
			return data;
		}

		public static void WriteReverse(this BinaryWriter writer, byte[] data) {
			var d = (byte[])data.Clone();
			Array.Reverse(d);
			writer.Write(d);
		}
	}

	/// <summary>
	/// A public key
	/// </summary>
	public sealed class StrongNamePublicKey {
		const uint RSA1_SIG = 0x31415352;
		SignatureAlgorithm signatureAlgorithm;
		AssemblyHashAlgorithm hashAlgorithm;
		byte[] modulus;
		byte[] publicExponent;

		/// <summary>
		/// Gets/sets the signature algorithm
		/// </summary>
		public SignatureAlgorithm SignatureAlgorithm {
			get { return signatureAlgorithm; }
			set { signatureAlgorithm = value; }
		}

		/// <summary>
		/// Gets/sets the hash algorithm
		/// </summary>
		public AssemblyHashAlgorithm HashAlgorithm {
			get { return hashAlgorithm; }
			set { hashAlgorithm = value; }
		}

		/// <summary>
		/// Gets/sets the modulus
		/// </summary>
		public byte[] Modulus {
			get { return modulus; }
			set { modulus = value; }
		}

		/// <summary>
		/// Gets/sets the public exponent
		/// </summary>
		public byte[] PublicExponent {
			get { return publicExponent; }
			set {
				if (publicExponent == null || publicExponent.Length != 4)
					throw new ArgumentException("PublicExponent must be exactly 4 bytes");
				publicExponent = value;
			}
		}

		/// <summary>
		/// Default constructor
		/// </summary>
		public StrongNamePublicKey() {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="modulus">Modulus</param>
		/// <param name="publicExponent">Public exponent</param>
		public StrongNamePublicKey(byte[] modulus, byte[] publicExponent)
			: this(modulus, publicExponent, AssemblyHashAlgorithm.SHA1, SignatureAlgorithm.CALG_RSA_SIGN) {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="modulus">Modulus</param>
		/// <param name="publicExponent">Public exponent</param>
		/// <param name="hashAlgorithm">Hash algorithm</param>
		public StrongNamePublicKey(byte[] modulus, byte[] publicExponent, AssemblyHashAlgorithm hashAlgorithm)
			: this(modulus, publicExponent, hashAlgorithm, SignatureAlgorithm.CALG_RSA_SIGN) {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="modulus">Modulus</param>
		/// <param name="publicExponent">Public exponent</param>
		/// <param name="hashAlgorithm">Hash algorithm</param>
		/// <param name="signatureAlgorithm">Signature algorithm</param>
		public StrongNamePublicKey(byte[] modulus, byte[] publicExponent, AssemblyHashAlgorithm hashAlgorithm, SignatureAlgorithm signatureAlgorithm) {
			this.signatureAlgorithm = signatureAlgorithm;
			this.hashAlgorithm = hashAlgorithm;
			this.modulus = modulus;
			this.publicExponent = publicExponent;
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="pk">Public key</param>
		public StrongNamePublicKey(PublicKey pk)
			: this(pk.Data) {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="pk">Public key data</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNamePublicKey(byte[] pk) {
			Initialize(new BinaryReader(new MemoryStream(pk)));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="filename">Public key file</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNamePublicKey(string filename) {
			using (var fileStream = File.OpenRead(filename))
				Initialize(new BinaryReader(fileStream));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="stream">Public key stream</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNamePublicKey(Stream stream) {
			Initialize(new BinaryReader(stream));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="reader">Public key reader</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNamePublicKey(BinaryReader reader) {
			Initialize(reader);
		}

		void Initialize(BinaryReader reader) {
			try {
				// Read PublicKeyBlob
				signatureAlgorithm = (SignatureAlgorithm)reader.ReadUInt32();
				hashAlgorithm = (AssemblyHashAlgorithm)reader.ReadUInt32();
				int pkLen = reader.ReadInt32();

				// Read PUBLICKEYSTRUC
				if (reader.ReadByte() != 6)
					throw new InvalidKeyException("Not a public key");
				if (reader.ReadByte() != 2)
					throw new InvalidKeyException("Invalid version");
				reader.ReadUInt16();	// reserved
				if ((SignatureAlgorithm)reader.ReadUInt32() != SignatureAlgorithm.CALG_RSA_SIGN)
					throw new InvalidKeyException("Not RSA sign");

				// Read RSAPUBKEY
				if (reader.ReadUInt32() != RSA1_SIG)	// magic = RSA1
					throw new InvalidKeyException("Invalid RSA1 magic");
				uint bitLength = reader.ReadUInt32();
				publicExponent = reader.ReadBytesReverse(4);

				modulus = reader.ReadBytesReverse((int)(bitLength / 8));
			}
			catch (IOException ex) {
				throw new InvalidKeyException("Invalid public key", ex);
			}
		}

		/// <summary>
		/// Creates a public key blob
		/// </summary>
		public byte[] CreatePublicKey() {
			return CreatePublicKey(signatureAlgorithm, hashAlgorithm, modulus, publicExponent);
		}

		internal static byte[] CreatePublicKey(SignatureAlgorithm sigAlg, AssemblyHashAlgorithm hashAlg, byte[] modulus, byte[] publicExponent) {
			if (sigAlg != SignatureAlgorithm.CALG_RSA_SIGN)
				throw new ArgumentException("Signature algorithm must be RSA");
			var outStream = new MemoryStream();
			var writer = new BinaryWriter(outStream);
			writer.Write((uint)sigAlg);		// SigAlgID
			writer.Write((uint)hashAlg);	// HashAlgID
			writer.Write(0x14 + modulus.Length);// cbPublicKey
			writer.Write((byte)6);			// bType (public key)
			writer.Write((byte)2);			// bVersion
			writer.Write((ushort)0);		// reserved
			writer.Write((uint)sigAlg);		// aiKeyAlg
			writer.Write(RSA1_SIG);			// magic (RSA1)
			writer.Write(modulus.Length * 8);	// bitlen
			writer.WriteReverse(publicExponent);// pubexp
			writer.WriteReverse(modulus);	// modulus
			return outStream.ToArray();
		}

		/// <inheritdoc/>
		public override string ToString() {
			return Utils.ToHex(CreatePublicKey(), false);
		}
	}

	/// <summary>
	/// Stores a strong name key pair
	/// </summary>
	public sealed class StrongNameKey {
		const uint RSA2_SIG = 0x32415352;
		byte[] publicKey;
		AssemblyHashAlgorithm hashAlg;
		byte[] publicExponent;
		byte[] modulus;
		byte[] prime1;
		byte[] prime2;
		byte[] exponent1;
		byte[] exponent2;
		byte[] coefficient;
		byte[] privateExponent;

		/// <summary>
		/// Gets the public key
		/// </summary>
		public byte[] PublicKey {
			get { return publicKey ?? (publicKey = CreatePublicKey()); }
		}

		/// <summary>
		/// Gets the strong name signature size in bytes
		/// </summary>
		public int SignatureSize {
			get { return modulus.Length; }
		}

		/// <summary>
		/// Gets/sets the public key hash algorithm. It's usually <see cref="AssemblyHashAlgorithm.SHA1"/>
		/// </summary>
		public AssemblyHashAlgorithm HashAlgorithm {
			get { return hashAlg; }
			set {
				if (hashAlg == value)
					return;
				publicKey = null;
				hashAlg = value;
			}
		}

		/// <summary>
		/// Gets/sets the public exponent
		/// </summary>
		public byte[] PublicExponent {
			get { return publicExponent; }
			set {
				if (publicExponent == null || publicExponent.Length != 4)
					throw new ArgumentException("PublicExponent must be exactly 4 bytes");
				publicExponent = value;
			}
		}

		/// <summary>
		/// Gets/sets the modulus
		/// </summary>
		public byte[] Modulus {
			get { return modulus; }
			set { modulus = value; }
		}

		/// <summary>
		/// Gets/sets prime1
		/// </summary>
		public byte[] Prime1 {
			get { return prime1; }
			set { prime1 = value; }
		}

		/// <summary>
		/// Gets/sets prime2
		/// </summary>
		public byte[] Prime2 {
			get { return prime2; }
			set { prime2 = value; }
		}

		/// <summary>
		/// Gets/sets exponent1
		/// </summary>
		public byte[] Exponent1 {
			get { return exponent1; }
			set { exponent1 = value; }
		}

		/// <summary>
		/// Gets/sets exponent2
		/// </summary>
		public byte[] Exponent2 {
			get { return exponent2; }
			set { exponent2 = value; }
		}

		/// <summary>
		/// Gets/sets the coefficient
		/// </summary>
		public byte[] Coefficient {
			get { return coefficient; }
			set { coefficient = value; }
		}

		/// <summary>
		/// Gets/sets the private exponent
		/// </summary>
		public byte[] PrivateExponent {
			get { return privateExponent; }
			set { privateExponent = value; }
		}

		/// <summary>
		/// Default constructor
		/// </summary>
		public StrongNameKey() {
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="keyData">Strong name key data</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNameKey(byte[] keyData) {
			Initialize(new BinaryReader(new MemoryStream(keyData)));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="filename">Strong name key file</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNameKey(string filename) {
			using (var fileStream = File.OpenRead(filename))
				Initialize(new BinaryReader(fileStream));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="stream">Strong name key stream</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNameKey(Stream stream) {
			Initialize(new BinaryReader(stream));
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="reader">Strong name key reader</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public StrongNameKey(BinaryReader reader) {
			Initialize(reader);
		}

		/// <summary>
		/// Initializes the public/private key pair data
		/// </summary>
		/// <param name="reader">Public/private key pair reader</param>
		/// <exception cref="InvalidKeyException">Strong name key is invalid</exception>
		public void Initialize(BinaryReader reader) {
			/*
			 * Links:
			 *	http://msdn.microsoft.com/en-us/library/cc250013%28v=prot.20%29.aspx
			 *	http://msdn.microsoft.com/en-us/library/windows/desktop/aa387689%28v=vs.85%29.aspx
			 *
			 *	struct PublicKeyBlob {
			 *	  unsigned int SigAlgID;	// sig algorithm used to create the sig (00002400 = CALG_RSA_SIGN)
			 *	  unsigned int HashAlgID;	// hash alg used to create the sig (usually 00008004 = CALG_SHA1)
			 *	  ULONG        cbPublicKey;	// Size of the data that follows
			 *	  // the rest is here
			 *	}
			 *
			 *	typedef struct _PUBLICKEYSTRUC {
			 *	  BYTE   bType;
			 *	  BYTE   bVersion;
			 *	  WORD   reserved;
			 *	  ALG_ID aiKeyAlg;
			 *	} BLOBHEADER, PUBLICKEYSTRUC;
			 *
			 *	typedef struct _RSAPUBKEY {
			 *	  DWORD magic;
			 *	  DWORD bitlen;
			 *	  DWORD pubexp;
			 *	} RSAPUBKEY;
			 *
			 * Format of public key
			 *	PublicKeyBlob
			 *	PUBLICKEYSTRUC	publickeystruc;
			 *	RSAPUBKEY		rsapubkey;
			 *	BYTE			modulus[rsapubkey.bitlen/8]
			 *
			 * Format of public/private key pair
			 *	PUBLICKEYSTRUC	publickeystruc;
			 *	RSAPUBKEY		rsapubkey;
			 *	BYTE			modulus[rsapubkey.bitlen/8];
			 *	BYTE			prime1[rsapubkey.bitlen/16];		// aka P
			 *	BYTE			prime2[rsapubkey.bitlen/16];		// aka Q
			 *	BYTE			exponent1[rsapubkey.bitlen/16];		// aka DP
			 *	BYTE			exponent2[rsapubkey.bitlen/16];		// aka DQ
			 *	BYTE			coefficient[rsapubkey.bitlen/16];	// aka IQ
			 *	BYTE			privateExponent[rsapubkey.bitlen/8];// aka D
			 */

			try {
				publicKey = null;

				// Read PUBLICKEYSTRUC
				if (reader.ReadByte() != 7)
					throw new InvalidKeyException("Not a public/private key pair");
				if (reader.ReadByte() != 2)
					throw new InvalidKeyException("Invalid version");
				reader.ReadUInt16();	// reserved
				if ((SignatureAlgorithm)reader.ReadUInt32() != SignatureAlgorithm.CALG_RSA_SIGN)
					throw new InvalidKeyException("Not RSA sign");

				// Read RSAPUBKEY
				if (reader.ReadUInt32() != RSA2_SIG)	// magic = RSA2
					throw new InvalidKeyException("Invalid RSA2 magic");
				uint bitLength = reader.ReadUInt32();
				publicExponent = reader.ReadBytesReverse(4);

				int len8 = (int)(bitLength / 8);
				int len16 = (int)(bitLength / 16);

				// Read the rest
				modulus = reader.ReadBytesReverse(len8);
				prime1 = reader.ReadBytesReverse(len16);
				prime2 = reader.ReadBytesReverse(len16);
				exponent1 = reader.ReadBytesReverse(len16);
				exponent2 = reader.ReadBytesReverse(len16);
				coefficient = reader.ReadBytesReverse(len16);
				privateExponent = reader.ReadBytesReverse(len8);
			}
			catch (IOException ex) {
				throw new InvalidKeyException("Couldn't read strong name key", ex);
			}
		}

		byte[] CreatePublicKey() {
			var halg = hashAlg == 0 ? AssemblyHashAlgorithm.SHA1 : hashAlg;
			return StrongNamePublicKey.CreatePublicKey(SignatureAlgorithm.CALG_RSA_SIGN, halg, modulus, publicExponent);
		}

		/// <summary>
		/// Creates an <see cref="RSA"/> instance
		/// </summary>
		public RSA CreateRSA() {
			var rsaParams = new RSAParameters {
				Exponent = publicExponent,
				Modulus = modulus,
				P = prime1,
				Q = prime2,
				DP = exponent1,
				DQ = exponent2,
				InverseQ = coefficient,
				D = privateExponent,
			};
			var rsa = RSA.Create();
			try {
				rsa.ImportParameters(rsaParams);
				return rsa;
			}
			catch {
				((IDisposable)rsa).Dispose();
				throw;
			}
		}

		/// <summary>
		/// Creates a strong name blob
		/// </summary>
		public byte[] CreateStrongName() {
			var outStream = new MemoryStream();
			var writer = new BinaryWriter(outStream);
			writer.Write((byte)7);			// bType (public/private key)
			writer.Write((byte)2);			// bVersion
			writer.Write((ushort)0);		// reserved
			writer.Write((uint)SignatureAlgorithm.CALG_RSA_SIGN);	// aiKeyAlg
			writer.Write(RSA2_SIG);			// magic (RSA2)
			writer.Write(modulus.Length * 8);	// bitlen
			writer.WriteReverse(publicExponent);
			writer.WriteReverse(modulus);
			writer.WriteReverse(prime1);
			writer.WriteReverse(prime2);
			writer.WriteReverse(exponent1);
			writer.WriteReverse(exponent2);
			writer.WriteReverse(coefficient);
			writer.WriteReverse(privateExponent);
			return outStream.ToArray();
		}

		/// <summary>
		/// Creates a counter signature, just like
		/// <c>sn -a IdentityPubKey.snk IdentityKey.snk SignaturePubKey.snk</c> can do.
		/// The public key <c>sn</c> prints is <paramref name="signaturePubKey"/>'s value.
		/// </summary>
		/// <param name="identityPubKey">Identity public key</param>
		/// <param name="identityKey">Identity strong name key pair</param>
		/// <param name="signaturePubKey">Signature public key</param>
		/// <returns>The counter signature as a hex string</returns>
		public static string CreateCounterSignatureAsString(StrongNamePublicKey identityPubKey, StrongNameKey identityKey, StrongNamePublicKey signaturePubKey) {
			var counterSignature = CreateCounterSignature(identityPubKey, identityKey, signaturePubKey);
			return Utils.ToHex(counterSignature, false);
		}

		/// <summary>
		/// Creates a counter signature, just like
		/// <c>sn -a IdentityPubKey.snk IdentityKey.snk SignaturePubKey.snk</c> can do.
		/// The public key <c>sn</c> prints is <paramref name="signaturePubKey"/>'s value.
		/// </summary>
		/// <param name="identityPubKey">Identity public key</param>
		/// <param name="identityKey">Identity strong name key pair</param>
		/// <param name="signaturePubKey">Signature public key</param>
		/// <returns>The counter signature</returns>
		public static byte[] CreateCounterSignature(StrongNamePublicKey identityPubKey, StrongNameKey identityKey, StrongNamePublicKey signaturePubKey) {
			var hash = AssemblyHash.Hash(signaturePubKey.CreatePublicKey(), identityPubKey.HashAlgorithm);
			using (var rsa = identityKey.CreateRSA()) {
				var rsaFmt = new RSAPKCS1SignatureFormatter(rsa);
				string hashName = identityPubKey.HashAlgorithm.GetName();
				rsaFmt.SetHashAlgorithm(hashName);
				var snSig = rsaFmt.CreateSignature(hash);
				Array.Reverse(snSig);
				return snSig;
			}
		}
	}
}
